import torch
import torch.nn as nn
from torch import optim
from torch.utils.tensorboard import SummaryWriter
#from MTL_trainfunction1 import train
from MTL_utils_1 import train

from MTL_dataloader1 import generate_train_data
from MTL_dataloader1 import generate_valid_data
from MTL_dataloader1 import generate_test_data
from MTL_SRmodel import MultiTaskSRNet


# device
device = torch.device('cuda:4' if torch.cuda.is_available() else 'cpu')

# 实例化模型
model = MultiTaskSRNet(data_format='NCHW', init_weights=True)
print('model down successfully')

# 数据预处理
#训练集使用多任务的训练集（隐写，滤波，滤波隐写），验证和测试的时候只使用隐写数据集
data_path = {
    'train_cover':'/data/user/mengxiangli/MTL_MMAL_dataset/BOWS2+Bossbase-cover/tc/',
    'train_stego': '/data/user/mengxiangli/MTL_MMAL_dataset/04bpp-WOW/ts/',
    'train_blurred':'/data/user/mengxiangli/MTL_MMAL_dataset/MTL_blurred/',
    'train_bs':'/data/user/mengxiangli/MTL_MMAL_dataset/MTL_bs_WOW_0.4/',

    'valid_cover': '/data/user/mengxiangli/MTL_MMAL_dataset/BOWS2+Bossbase-cover/vc/',
    'valid_stego': '/data/user/mengxiangli/MTL_MMAL_dataset/04bpp-WOW/vs/',
    
    'test_cover': '/data/user/mengxiangli/MTL_MMAL_dataset/BOWS2+Bossbase-cover/testc/',
    'test_stego': '/data/user/mengxiangli/MTL_MMAL_dataset/04bpp-WOW/tests/'
}


batch_size = {'train': 16, 'valid': 16,'test':16}
train_loader = generate_train_data(data_path, batch_size)
valid_loader = generate_valid_data(data_path, batch_size)
test_loader = generate_test_data(data_path, batch_size)
print('data_loader down successfully')

# 训练参数设置
EPOCHS = 180
write_interval = 100
valid_interval = 875
save_interval = 10000
learning_rate = 1e-3

# 损失函数 优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
load_path = None
#load_path = './Mo_3x3_0.6bpp_HUGO_best/bestModel_42000_0.9227.pth'
# 开始训练
print('start train')
train(model,
    train_loader,
    valid_loader,
    test_loader,
    optimizer,
    criterion,
    device,
    EPOCHS,
    save_interval=5,
    lr_adjust_epochs=[80, 120],
    gamma=0.1,
    early_stop_patience=25,
    early_stop_threshold=0.85,
    log_dir='runs',
    checkpoint_path=None)

